{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# Anatomy of a Quantizer\n",
    "\n",
    "## What's in a Quantizer? \n",
    "\n",
    "In a broad sense, a quantizer is anything that implements a quantization technique, and the flexibility of Brevitas means that there are different ways to do so.  \n",
    "However, to keep our terminology straight, we refer to a quantizer as a specific kind of way to implement quantization, the one preferred and adopted by default. That is, a quantizer is a subclass of a `brevitas.inject.ExtendedInjector` that carries a `tensor_quant` attribute, which points to an instance of a torch `Module` that implements quantization.\n",
    "\n",
    "We have seen in previous tutorials quantizers being imported from `brevitas.quant` and passed on to quantized layers. We can easily very what we just said on one of them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:29.666286Z",
     "iopub.status.busy": "2025-05-09T14:48:29.665562Z",
     "iopub.status.idle": "2025-05-09T14:48:36.527900Z",
     "shell.execute_reply": "2025-05-09T14:48:36.525705Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.inject import ExtendedInjector\n",
    "from brevitas.quant.scaled_int import Int8ActPerTensorFloat \n",
    "\n",
    "issubclass(Int8ActPerTensorFloat, ExtendedInjector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:36.672297Z",
     "iopub.status.busy": "2025-05-09T14:48:36.671220Z",
     "iopub.status.idle": "2025-05-09T14:48:37.115201Z",
     "shell.execute_reply": "2025-05-09T14:48:37.112900Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RescalingIntQuant(\n",
       "  (int_quant): IntQuant(\n",
       "    (float_to_int_impl): RoundSte()\n",
       "    (tensor_clamp_impl): TensorClamp()\n",
       "    (delay_wrapper): DelayWrapper(\n",
       "      (delay_impl): _NoDelay()\n",
       "    )\n",
       "    (input_view_impl): Identity()\n",
       "  )\n",
       "  (scaling_impl): ParameterFromRuntimeStatsScaling(\n",
       "    (stats_input_view_shape_impl): OverTensorView()\n",
       "    (stats): _Stats(\n",
       "      (stats_impl): AbsPercentile()\n",
       "    )\n",
       "    (restrict_scaling): _RestrictValue(\n",
       "      (restrict_value_impl): FloatRestrictValue()\n",
       "    )\n",
       "    (restrict_threshold): _RestrictValue(\n",
       "      (restrict_value_impl): FloatRestrictValue()\n",
       "    )\n",
       "    (clamp_scaling): _ClampValue(\n",
       "      (clamp_min_ste): ScalarClampMinSte()\n",
       "    )\n",
       "    (restrict_inplace_preprocess): Identity()\n",
       "    (restrict_scaling_pre): Identity()\n",
       "    (restrict_threshold_pre): Identity()\n",
       "  )\n",
       "  (int_scaling_impl): IntScaling()\n",
       "  (zero_point_impl): ZeroZeroPoint(\n",
       "    (zero_point): StatelessBuffer()\n",
       "  )\n",
       "  (msb_clamp_bit_width_impl): BitWidthConst(\n",
       "    (bit_width): StatelessBuffer()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Int8ActPerTensorFloat.tensor_quant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how we said *subclass* and not *instance*. To understand why that's the case, we have to understand what an `ExtendedInjector` is and why it's used in the first place."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization with auto-wiring Dependency Injection\n",
    "\n",
    "Pytorch has exploded in popularity thanks to its straightforward numpy-like *define-by-run* execution model. However, when it comes to applying quantization, this style of programming poses a problem.  \n",
    "\n",
    "Many quantization methods depend on making decisions based on the (in Pytorch terms) `state_dict` of the original floating-point model to finetune with quantization. However, when we instantiate a model in Pytorch we can't know on the spot if a state_dict is going to be loaded a few lines of code later or not. Yet, because Pytorch is define-by-run, we need our model to work consistently both before and after a `state_dict` is possibly loaded. In a traditional scenario that wouldn't pose a problem. However, with quantization in the loop, the way a quantizer is defined might change before and after a pretrained `state_dict` is loaded.\n",
    "\n",
    "That means that we need a way to define our quantized model such that it can react appropriately in case the `state_dict` changes. In a Python-only world that wouldn't be too hard. However, in order to mitigate the performance impact of quantization-aware training, Brevitas makes extended use of Pytorch's JIT compiler for a custom subset of Python, TorchScript. That means that in most scenarios, when a `state_dict` is loaded, we need to recompile parts of the model. Because compilation in general is a lossy process, a TorchScript component cannot simply re-compile itself based on new input information. \n",
    "\n",
    "We need then a way to *declare* a quantization method such that it can be re-initialized and JIT compiled any time the `state_dict` changes. Because we want to support arbitrarly-complex user-defined quantization algorithms, this method has to be generic, i.e. it cannot depend on the specifics of the quantization algorithm implemented.\n",
    "\n",
    "Implementing a quantizer with an `ExtendedInjector` is a way to do so. Specifically, an `ExtendedInjector` extends an `Injector` from an older version (*0.2.1*) of the excellent dependency-injection library *[dependencies](https://github.com/proofit404/dependencies)* with support for a couple of extra features that are specific to Brevitas' needs.\n",
    "\n",
    "An `Injector` (and an `ExtendedInjector`) allows to take what might be a very complicated graph of interwined objects and turns it into a flat list of variables that are capable of auto-assembly by matching variable names to arguments names. This technique typically goes under the name of auto-wiring dependency injection. \n",
    "\n",
    "In the context of Brevitas, the goal is gather all the modules and hyperparameters that contribute to a quantization implementation such that they can be re-assembled automatically on demand. What comes out of this process is a `tensor_quant` object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Practical Example: Binary Quantization\n",
    "\n",
    "To make things practical, let's look at how we can implement a simple variant of binary quantization. All the components typically used to implement quantization can be found under `brevitas.core`. As mentioned before, Brevitas makes heavy use of TorchScript. In particular, all the components found under `brevitas.core` are implemented as `ScriptModule` that can be assembled together.\n",
    "The core `ScriptModule` that implements binarization can be found under `brevitas.core.quant`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.128784Z",
     "iopub.status.busy": "2025-05-09T14:48:37.128045Z",
     "iopub.status.idle": "2025-05-09T14:48:37.141243Z",
     "shell.execute_reply": "2025-05-09T14:48:37.139062Z"
    }
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "# helpers\n",
    "def assert_with_message(condition):\n",
    "    assert condition\n",
    "    print(condition)\n",
    "\n",
    "def pretty_print_source(source):\n",
    "    display(Markdown('```python\\n' + source + '\\n```'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.153497Z",
     "iopub.status.busy": "2025-05-09T14:48:37.152581Z",
     "iopub.status.idle": "2025-05-09T14:48:37.177029Z",
     "shell.execute_reply": "2025-05-09T14:48:37.174495Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "class BinaryQuant(brevitas.jit.ScriptModule):\n",
       "    \"\"\"\n",
       "    ScriptModule that implements scaled uniform binary quantization of an input tensor.\n",
       "    Quantization is performed with :func:`~brevitas.function.ops_ste.binary_sign_ste`.\n",
       "\n",
       "    Args:\n",
       "        scaling_impl (Module): Module that returns a scale factor.\n",
       "        quant_delay_steps (int): Number of training steps to delay quantization for. Default: 0\n",
       "\n",
       "    Returns:\n",
       "        Tuple[Tensor, Tensor, Tensor, Tensor]: Quantized output in de-quantized format, scale, zero-point, bit_width.\n",
       "\n",
       "    Examples:\n",
       "        >>> from brevitas.core.scaling import ConstScaling\n",
       "        >>> binary_quant = BinaryQuant(ConstScaling(0.1))\n",
       "        >>> inp = torch.Tensor([0.04, -0.6, 3.3])\n",
       "        >>> out, scale, zero_point, bit_width = binary_quant(inp)\n",
       "        >>> out\n",
       "        tensor([ 0.1000, -0.1000,  0.1000])\n",
       "        >>> scale\n",
       "        tensor(0.1000)\n",
       "        >>> zero_point\n",
       "        tensor(0.)\n",
       "        >>> bit_width\n",
       "        tensor(1.)\n",
       "\n",
       "    Note:\n",
       "        Maps to quant_type == QuantType.BINARY == 'BINARY' == 'binary' when applied to weights in higher-level APIs.\n",
       "\n",
       "    Note:\n",
       "        Set env variable BREVITAS_JIT=1 to enable TorchScript compilation of this module.\n",
       "    \"\"\"\n",
       "\n",
       "    def __init__(self, scaling_impl: Module, signed: bool = True, quant_delay_steps: int = 0):\n",
       "        super(BinaryQuant, self).__init__()\n",
       "        assert signed, \"Unsigned binary quant not supported\"\n",
       "        self.scaling_impl = scaling_impl\n",
       "        self.bit_width = BitWidthConst(1)\n",
       "        self.zero_point = StatelessBuffer(torch.tensor(0.0))\n",
       "        self.delay_wrapper = DelayWrapper(quant_delay_steps)\n",
       "\n",
       "    @brevitas.jit.script_method\n",
       "    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:\n",
       "        scale = self.scaling_impl(x)\n",
       "        y = binary_sign_ste(x) * scale\n",
       "        y = self.delay_wrapper(x, y)\n",
       "        return y, scale, self.zero_point(), self.bit_width()\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from brevitas.core.quant import BinaryQuant\n",
    "\n",
    "source = inspect.getsource(BinaryQuant)  \n",
    "pretty_print_source(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The implementation is quite simple. Apart from `quant_delay_steps`, which allows to delay quantization by a certain number of training steps (*default = 0*), the only other argument that BinaryQuant accepts is an implementation to compute the scale factor. `bit_width` is fixed to 1 and `zero-point` is fixed to 0.\n",
    "\n",
    "We pick as scale factor implementation a `ScriptModule` called `ParameterScaling`, which implements a learned parameter with user-defined initialization. It can be found under `brevitas.core.scaling`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.186317Z",
     "iopub.status.busy": "2025-05-09T14:48:37.185624Z",
     "iopub.status.idle": "2025-05-09T14:48:37.195320Z",
     "shell.execute_reply": "2025-05-09T14:48:37.193140Z"
    }
   },
   "outputs": [],
   "source": [
    "from brevitas.core.scaling import ParameterScaling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manual Binary Quantization\n",
    "\n",
    "As a first step, we simply instantiate `BinaryQuant` with `ParameterScaling` using `scaling_init` equal *0.1* and we call it on a random floating-point input tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.204168Z",
     "iopub.status.busy": "2025-05-09T14:48:37.203457Z",
     "iopub.status.idle": "2025-05-09T14:48:37.236183Z",
     "shell.execute_reply": "2025-05-09T14:48:37.233948Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.1000, -0.1000, -0.1000, -0.1000],\n",
       "         [ 0.1000,  0.1000, -0.1000, -0.1000],\n",
       "         [ 0.1000, -0.1000,  0.1000,  0.1000],\n",
       "         [ 0.1000,  0.1000,  0.1000, -0.1000]], grad_fn=<MulBackward0>),\n",
       " tensor([0.1000], grad_fn=<DivBackward0>),\n",
       " tensor(0.),\n",
       " tensor(1.))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "# set seed for notebook\n",
    "torch.manual_seed(0)\n",
    "\n",
    "manual_tensor_quant = BinaryQuant(scaling_impl=ParameterScaling(scaling_init=0.1))\n",
    "manual_tensor_quant(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nothing too surprising here, as expected the tensor is binarized with the scale factor we defined. Note however how `manual_tensor_quant` is returning a `tuple` and not a `QuantTensor`. This is because support for custom data structures in TorchScript is still quite limited, so `QuantTensor` are allocated only in Python-world abstractions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binary Quantization with an ExtendedInjector\n",
    "\n",
    "Let's now declare `tensor_quant` through an `ExtendedInjector`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.246429Z",
     "iopub.status.busy": "2025-05-09T14:48:37.245639Z",
     "iopub.status.idle": "2025-05-09T14:48:37.268642Z",
     "shell.execute_reply": "2025-05-09T14:48:37.266398Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.1000, -0.1000,  0.1000,  0.1000],\n",
       "         [ 0.1000, -0.1000, -0.1000,  0.1000],\n",
       "         [ 0.1000, -0.1000, -0.1000,  0.1000],\n",
       "         [ 0.1000,  0.1000,  0.1000, -0.1000]], grad_fn=<MulBackward0>),\n",
       " tensor([0.1000], grad_fn=<DivBackward0>),\n",
       " tensor(0.),\n",
       " tensor(1.))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.inject import ExtendedInjector\n",
    "\n",
    "class MyBinaryQuantizer(ExtendedInjector):\n",
    "    tensor_quant = BinaryQuant\n",
    "    scaling_impl=ParameterScaling\n",
    "    scaling_init=0.1\n",
    "\n",
    "inj_tensor_quant = MyBinaryQuantizer.tensor_quant\n",
    "inj_tensor_quant(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Any time `MyBinaryQuantizer.tensor_quant` is called, a new instance of `BinaryQuant` is created. Note how the attributes of `MyBinaryQuantizer` are designed to match the name of the arguments of each other, except for `tensor_quant`, which is what we are interested in retrieving from the outside."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inheritance and Composition of Quantizers\n",
    "\n",
    "The advantage of expressing a quantizer through a Python class also means that we can leverage both *inheritance* and *composition*. So for example we can inherit from `MyBinaryQuantizer` and override `scaling_init` with a new value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.278274Z",
     "iopub.status.busy": "2025-05-09T14:48:37.277458Z",
     "iopub.status.idle": "2025-05-09T14:48:37.299093Z",
     "shell.execute_reply": "2025-05-09T14:48:37.296827Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-1.,  1., -1.,  1.],\n",
       "         [ 1.,  1.,  1.,  1.],\n",
       "         [-1.,  1., -1.,  1.],\n",
       "         [ 1.,  1., -1., -1.]], grad_fn=<MulBackward0>),\n",
       " tensor([1.], grad_fn=<DivBackward0>),\n",
       " tensor(0.),\n",
       " tensor(1.))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class MyChildBinaryQuantizer(MyBinaryQuantizer):\n",
    "    scaling_init=1.0\n",
    "    \n",
    "child_inj_tensor_quant = MyChildBinaryQuantizer.tensor_quant\n",
    "child_inj_tensor_quant(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or we can leverage composition by assembling together various classes containing different pieces of a quantizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.308050Z",
     "iopub.status.busy": "2025-05-09T14:48:37.307351Z",
     "iopub.status.idle": "2025-05-09T14:48:37.331563Z",
     "shell.execute_reply": "2025-05-09T14:48:37.329296Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.1000,  0.1000,  0.1000,  0.1000],\n",
       "         [-0.1000, -0.1000,  0.1000, -0.1000],\n",
       "         [-0.1000,  0.1000,  0.1000,  0.1000],\n",
       "         [-0.1000, -0.1000,  0.1000, -0.1000]], grad_fn=<MulBackward0>),\n",
       " tensor([0.1000], grad_fn=<DivBackward0>),\n",
       " tensor(0.),\n",
       " tensor(1.))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class MyBinaryImpl(ExtendedInjector):\n",
    "    tensor_quant = BinaryQuant\n",
    "\n",
    "class MyScalingImpl(ExtendedInjector):\n",
    "    scaling_impl=ParameterScaling\n",
    "    scaling_init=0.1\n",
    "    \n",
    "class MyComposedBinaryQuantizer(MyBinaryImpl, MyScalingImpl):\n",
    "    pass\n",
    "\n",
    "comp_inj_tensor_quant = MyComposedBinaryQuantizer.tensor_quant\n",
    "comp_inj_tensor_quant(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interfacing a Quantizer with a Quantized Layer\n",
    "\n",
    "Before we can pass the quantizer to a quantized layer such as `QuantConv2d`, we need a last component to define, a proxy. A proxy (found under `brevitas.proxy`) is an `nn.Module` that serve as interface between a quantizer and a quantized layer. \n",
    "\n",
    "While a quantizer lives mostly in JIT-land, a proxy lives mostly in Python-land, and as such can afford much more flexibility. Proxies take care of returning a `QuantTensor` and re-initializing the output of quantizer whenever a new `state_dict` is loaded. \n",
    "\n",
    "Proxies are specific to the kind of tensor being quantized, as in weights vs biases vs activations. For convenience, they are declared as part of the quantizer itself under the attribute `proxy_class`. For example, for weights we can use `WeightQuantProxyFromInjector`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.340768Z",
     "iopub.status.busy": "2025-05-09T14:48:37.340010Z",
     "iopub.status.idle": "2025-05-09T14:48:37.350031Z",
     "shell.execute_reply": "2025-05-09T14:48:37.347971Z"
    }
   },
   "outputs": [],
   "source": [
    "from brevitas.proxy import WeightQuantProxyFromInjector\n",
    "\n",
    "class MyBinaryWeightQuantizer(MyBinaryQuantizer):\n",
    "    proxy_class = WeightQuantProxyFromInjector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use `MyBinaryWeightQuantizer` as the weight quantizer of a layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.359031Z",
     "iopub.status.busy": "2025-05-09T14:48:37.358342Z",
     "iopub.status.idle": "2025-05-09T14:48:37.441323Z",
     "shell.execute_reply": "2025-05-09T14:48:37.439478Z"
    }
   },
   "outputs": [],
   "source": [
    "from brevitas.nn import QuantConv2d\n",
    "\n",
    "binary_weight_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MyBinaryWeightQuantizer)\n",
    "try:\n",
    "    quant_weight = binary_weight_quant_conv.quant_weight()\n",
    "except TypeError:\n",
    "    pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note however that we cannot compute the quantized weight, as the `signed` attribute is `None`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`signed` is one of those attributes that in the case of binary quantization has to be explicitly defined by the user. The idea is that it informs the proxy on whether the value generated by our quantizer should be considered signed or not. We can do so by simply setting it in the quantizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.451657Z",
     "iopub.status.busy": "2025-05-09T14:48:37.450761Z",
     "iopub.status.idle": "2025-05-09T14:48:37.482020Z",
     "shell.execute_reply": "2025-05-09T14:48:37.479869Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[ 0.1000,  0.1000, -0.1000],\n",
       "          [ 0.1000, -0.1000, -0.1000],\n",
       "          [-0.1000, -0.1000, -0.1000]],\n",
       "\n",
       "         [[-0.1000, -0.1000,  0.1000],\n",
       "          [-0.1000, -0.1000,  0.1000],\n",
       "          [-0.1000,  0.1000, -0.1000]],\n",
       "\n",
       "         [[ 0.1000,  0.1000,  0.1000],\n",
       "          [ 0.1000,  0.1000, -0.1000],\n",
       "          [-0.1000, -0.1000,  0.1000]]],\n",
       "\n",
       "\n",
       "        [[[-0.1000, -0.1000, -0.1000],\n",
       "          [-0.1000, -0.1000, -0.1000],\n",
       "          [ 0.1000,  0.1000, -0.1000]],\n",
       "\n",
       "         [[ 0.1000, -0.1000, -0.1000],\n",
       "          [-0.1000, -0.1000, -0.1000],\n",
       "          [-0.1000, -0.1000, -0.1000]],\n",
       "\n",
       "         [[ 0.1000, -0.1000,  0.1000],\n",
       "          [-0.1000, -0.1000,  0.1000],\n",
       "          [-0.1000,  0.1000,  0.1000]]]], grad_fn=<MulBackward0>), scale=tensor([0.1000], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class MySignedBinaryWeightQuantizer(MyBinaryWeightQuantizer):\n",
    "    signed = True\n",
    "    \n",
    "binary_weight_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MySignedBinaryWeightQuantizer)\n",
    "signed_quant_weight = binary_weight_quant_conv.quant_weight()\n",
    "signed_quant_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.490322Z",
     "iopub.status.busy": "2025-05-09T14:48:37.489641Z",
     "iopub.status.idle": "2025-05-09T14:48:37.501678Z",
     "shell.execute_reply": "2025-05-09T14:48:37.500153Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message(signed_quant_weight.is_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now the quant weights are valid.\n",
    "\n",
    "When we want to add or override an single attribute of a quantizer passed to a layer, defining a whole new quantizer can be too verbose. There is a simpler syntax to achieve the same goal. Let's say we want to have add the `signed` attribute to `MyBinaryQuantizer`, as we just did. We could have also simply done the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.510483Z",
     "iopub.status.busy": "2025-05-09T14:48:37.509806Z",
     "iopub.status.idle": "2025-05-09T14:48:37.539269Z",
     "shell.execute_reply": "2025-05-09T14:48:37.537245Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[-0.1000, -0.1000,  0.1000],\n",
       "          [-0.1000,  0.1000, -0.1000],\n",
       "          [ 0.1000,  0.1000,  0.1000]],\n",
       "\n",
       "         [[ 0.1000, -0.1000, -0.1000],\n",
       "          [-0.1000,  0.1000,  0.1000],\n",
       "          [ 0.1000, -0.1000,  0.1000]],\n",
       "\n",
       "         [[-0.1000, -0.1000, -0.1000],\n",
       "          [ 0.1000, -0.1000, -0.1000],\n",
       "          [ 0.1000,  0.1000, -0.1000]]],\n",
       "\n",
       "\n",
       "        [[[-0.1000, -0.1000,  0.1000],\n",
       "          [-0.1000,  0.1000,  0.1000],\n",
       "          [ 0.1000, -0.1000, -0.1000]],\n",
       "\n",
       "         [[ 0.1000, -0.1000, -0.1000],\n",
       "          [ 0.1000, -0.1000, -0.1000],\n",
       "          [ 0.1000, -0.1000, -0.1000]],\n",
       "\n",
       "         [[ 0.1000,  0.1000,  0.1000],\n",
       "          [-0.1000,  0.1000, -0.1000],\n",
       "          [-0.1000, -0.1000, -0.1000]]]], grad_fn=<MulBackward0>), scale=tensor([0.1000], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_scale_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MyBinaryWeightQuantizer, weight_signed=True)\n",
    "small_scale_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What we did was to take the name of the attribute `signed`, add the prefix `weight_`, and pass it as a keyword argument to `QuantConv2d`. What happens in the background is that the keyword arguments prefixed with `weight_` are set as attributes of `weight_quant`, possibly overriding any pre-existing value. The same principle applies to `input_`, `output_` and `bias_`.\n",
    "\n",
    "This is the reason why, as it was mentioned in the first tutorial, quantized layers can accept arbitrary keyword arguments. It's really just a way to support different styles of syntax when defining a quantizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Passing a custom quantizer to QuantIdentity\n",
    "\n",
    "We can do a similar thing with quantized activations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.548416Z",
     "iopub.status.busy": "2025-05-09T14:48:37.547737Z",
     "iopub.status.idle": "2025-05-09T14:48:37.575653Z",
     "shell.execute_reply": "2025-05-09T14:48:37.573780Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/_tensor.py:1255: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /opt/conda/conda-bld/pytorch_1670525541990/work/c10/core/TensorImpl.h:1758.)\n",
      "  return super(Tensor, self).rename(names)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[-0.1000,  0.1000, -0.1000, -0.1000],\n",
       "        [ 0.1000,  0.1000,  0.1000,  0.1000],\n",
       "        [-0.1000, -0.1000,  0.1000, -0.1000],\n",
       "        [-0.1000,  0.1000, -0.1000,  0.1000]], grad_fn=<MulBackward0>), scale=tensor([0.1000], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.proxy import ActQuantProxyFromInjector\n",
    "from brevitas.nn import QuantIdentity\n",
    "\n",
    "class MySignedBinaryActQuantizer(MyBinaryQuantizer):\n",
    "    proxy_class = ActQuantProxyFromInjector\n",
    "    signed = True\n",
    "\n",
    "binary_relu = QuantIdentity(act_quant=MySignedBinaryActQuantizer, return_quant_tensor=True)\n",
    "binary_relu(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So there isn't really much difference between a quantizer for weights and a quantizer for activations, they are just wrapped by different proxies. Also, with activations  a prefix is not required when passing keyword arguments. For example, when can override the existing `scaling_init` defined in `MyBinaryQuantizer` with a new value passed in as a keywork argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.584524Z",
     "iopub.status.busy": "2025-05-09T14:48:37.583707Z",
     "iopub.status.idle": "2025-05-09T14:48:37.607234Z",
     "shell.execute_reply": "2025-05-09T14:48:37.605227Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[ 0.0010, -0.0010, -0.0010,  0.0010],\n",
       "        [ 0.0010,  0.0010,  0.0010,  0.0010],\n",
       "        [ 0.0010, -0.0010,  0.0010,  0.0010],\n",
       "        [ 0.0010, -0.0010, -0.0010, -0.0010]], grad_fn=<MulBackward0>), scale=tensor([0.0010], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_scale_binary_identity = QuantIdentity(\n",
    "    act_quant=MySignedBinaryActQuantizer, scaling_init=0.001, return_quant_tensor=True)\n",
    "small_scale_binary_identity(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Custom Quantizer initialized with Weight Statistics\n",
    "\n",
    "So far we have seen use-cases where an `ExtendedInjector` provides, at best, a different kind of syntax to define a quantizer, without any particular other advantage. Let's now make things a bit more complicated to show the sort of situations where it really shines.\n",
    "\n",
    "Let's say we want to define a binary weight quantizer where `scaling_impl` is still `ParameterScaling`. However, instead of being user-defined, we want `scaling_init` to be the maximum value found in the weight tensor of the quantized layer.\n",
    "To support this sort of use cases where the quantizer depends on the layer, a quantized layer automatically passes itself to all its quantizers under the name of `module`.\n",
    "With only a few lines of code then, we can achieve our goal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.615921Z",
     "iopub.status.busy": "2025-05-09T14:48:37.615218Z",
     "iopub.status.idle": "2025-05-09T14:48:37.626346Z",
     "shell.execute_reply": "2025-05-09T14:48:37.624318Z"
    }
   },
   "outputs": [],
   "source": [
    "from brevitas.inject import value\n",
    "\n",
    "class ParamFromMaxWeightQuantizer(MySignedBinaryWeightQuantizer):\n",
    "    \n",
    "    @value\n",
    "    def scaling_init(module):\n",
    "        return module.weight.abs().max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how we are leveraging the `@value` *decorator* to define a function that is executed at *dependency-injection (DI) time*. This kind of behaviour is similar in spirit to defining a `@property` instead of an *attribute*, with the difference that a `@value` function can depend on other attributes of the Injector, which are automatically passed in as arguments of the function during DI.\n",
    "\n",
    "Let's now pass the quantizer to a QuantConv2d and retrieve its quantized weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.635130Z",
     "iopub.status.busy": "2025-05-09T14:48:37.634429Z",
     "iopub.status.idle": "2025-05-09T14:48:37.664437Z",
     "shell.execute_reply": "2025-05-09T14:48:37.662394Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[ 0.1820, -0.1820, -0.1820],\n",
       "          [ 0.1820,  0.1820,  0.1820],\n",
       "          [-0.1820, -0.1820, -0.1820]],\n",
       "\n",
       "         [[ 0.1820, -0.1820, -0.1820],\n",
       "          [ 0.1820, -0.1820, -0.1820],\n",
       "          [ 0.1820,  0.1820, -0.1820]],\n",
       "\n",
       "         [[-0.1820,  0.1820,  0.1820],\n",
       "          [ 0.1820, -0.1820,  0.1820],\n",
       "          [-0.1820, -0.1820, -0.1820]]],\n",
       "\n",
       "\n",
       "        [[[ 0.1820,  0.1820, -0.1820],\n",
       "          [-0.1820, -0.1820,  0.1820],\n",
       "          [-0.1820,  0.1820, -0.1820]],\n",
       "\n",
       "         [[-0.1820,  0.1820, -0.1820],\n",
       "          [ 0.1820,  0.1820, -0.1820],\n",
       "          [ 0.1820,  0.1820,  0.1820]],\n",
       "\n",
       "         [[-0.1820, -0.1820, -0.1820],\n",
       "          [ 0.1820, -0.1820,  0.1820],\n",
       "          [ 0.1820, -0.1820,  0.1820]]]], grad_fn=<MulBackward0>), scale=tensor([0.1820], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "param_from_max_quant_conv = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)\n",
    "param_from_max_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed we can verify that `weight_quant.scale()` is equal to `weight.abs().max()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.673587Z",
     "iopub.status.busy": "2025-05-09T14:48:37.672907Z",
     "iopub.status.idle": "2025-05-09T14:48:37.684677Z",
     "shell.execute_reply": "2025-05-09T14:48:37.682851Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message((param_from_max_quant_conv.weight_quant.scale() == param_from_max_quant_conv.weight.abs().max()).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say now that we want to load a pretrained floating-point weight tensor on top of our quantized model. We simuate this scenario by defining a separate `nn.Conv2d` layer with the same weight shape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.693605Z",
     "iopub.status.busy": "2025-05-09T14:48:37.692926Z",
     "iopub.status.idle": "2025-05-09T14:48:37.708605Z",
     "shell.execute_reply": "2025-05-09T14:48:37.706772Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.1924, grad_fn=<MaxBackward1>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch import nn\n",
    "\n",
    "float_conv = nn.Conv2d(3, 2, (3, 3))\n",
    "float_conv.weight.abs().max()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and then we load it on top of `param_from_max_quant_conv`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:37.717550Z",
     "iopub.status.busy": "2025-05-09T14:48:37.716868Z",
     "iopub.status.idle": "2025-05-09T14:48:40.330474Z",
     "shell.execute_reply": "2025-05-09T14:48:40.327964Z"
    },
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Error(s) in loading state_dict for QuantConv2d:\n\tMissing key(s) in state_dict: \"weight_quant.tensor_quant.scaling_impl.value\". ",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[21], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mparam_from_max_quant_conv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_state_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfloat_conv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstate_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py:1671\u001b[0m, in \u001b[0;36mModule.load_state_dict\u001b[0;34m(self, state_dict, strict)\u001b[0m\n\u001b[1;32m   1666\u001b[0m         error_msgs\u001b[38;5;241m.\u001b[39minsert(\n\u001b[1;32m   1667\u001b[0m             \u001b[38;5;241m0\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mMissing key(s) in state_dict: \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m. \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m   1668\u001b[0m                 \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m, \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(k) \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m missing_keys)))\n\u001b[1;32m   1670\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(error_msgs) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m-> 1671\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mError(s) in loading state_dict for \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m:\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m   1672\u001b[0m                        \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(error_msgs)))\n\u001b[1;32m   1673\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _IncompatibleKeys(missing_keys, unexpected_keys)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Error(s) in loading state_dict for QuantConv2d:\n\tMissing key(s) in state_dict: \"weight_quant.tensor_quant.scaling_impl.value\". "
     ]
    }
   ],
   "source": [
    "param_from_max_quant_conv.load_state_dict(float_conv.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ouch, we get an error. This is because `ParameterScaling` contains a learned `torch.nn.Parameter`, and Pytorch expects all learned parameters of a model to be contained in a `state_dict` that is being loaded.\n",
    "We can work around the issue by either setting the `IGNORE_MISSING_KEYS` config flag in Brevitas, or by passing `strict=False` to load_state_dict. We go with the former as setting `strict=False` is too forgiving to other kind of problems:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.341814Z",
     "iopub.status.busy": "2025-05-09T14:48:40.340146Z",
     "iopub.status.idle": "2025-05-09T14:48:40.361525Z",
     "shell.execute_reply": "2025-05-09T14:48:40.359798Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas import config\n",
    "config.IGNORE_MISSING_KEYS = True\n",
    "\n",
    "param_from_max_quant_conv.load_state_dict(float_conv.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we could have also achieve the same goal by setting the *env variable* `BREVITAS_IGNORE_MISSING_KEYS=1`.\n",
    "\n",
    "And now if we take a look at the quantized weights again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.371940Z",
     "iopub.status.busy": "2025-05-09T14:48:40.370319Z",
     "iopub.status.idle": "2025-05-09T14:48:40.394793Z",
     "shell.execute_reply": "2025-05-09T14:48:40.392661Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[ 0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924]],\n",
       "\n",
       "         [[-0.1924, -0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924]],\n",
       "\n",
       "         [[ 0.1924,  0.1924, -0.1924],\n",
       "          [-0.1924,  0.1924,  0.1924],\n",
       "          [-0.1924,  0.1924,  0.1924]]],\n",
       "\n",
       "\n",
       "        [[[ 0.1924, -0.1924,  0.1924],\n",
       "          [-0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924]],\n",
       "\n",
       "         [[-0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924, -0.1924, -0.1924],\n",
       "          [-0.1924,  0.1924,  0.1924]],\n",
       "\n",
       "         [[ 0.1924,  0.1924, -0.1924],\n",
       "          [-0.1924, -0.1924, -0.1924],\n",
       "          [ 0.1924, -0.1924, -0.1924]]]], grad_fn=<MulBackward0>), scale=tensor([0.1924], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "param_from_max_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that, as expected, the scale factor has been updated to the new `weight.abs().max()`.\n",
    "\n",
    "What happens internally is that after `load_state_dict` is called on the layer, `ParamFromMaxWeightQuantizer.tensor_quant` gets called again to re-initialize `BinaryQuant`, and in turn `ParameterScaling` is re-initialized with a new `scaling_init` value computed based on the updated `module.weight` tensor. This whole process wouldn't have been possible without an `ExtendedInjector` behind it. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sharing a Quantizer\n",
    "\n",
    "There are two ways to share a quantizer between multiple layers, with importance differences.\n",
    "\n",
    "The first one, which we have seen so far, is to simply pass the same ExtendedInjector to multiple layers. What that does is sharing the same quantization strategy among different layers. Each layer still gets its own instance of the quantization implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.405120Z",
     "iopub.status.busy": "2025-05-09T14:48:40.403554Z",
     "iopub.status.idle": "2025-05-09T14:48:40.436463Z",
     "shell.execute_reply": "2025-05-09T14:48:40.434140Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=MySignedBinaryWeightQuantizer)\n",
    "quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=MySignedBinaryWeightQuantizer)\n",
    "\n",
    "assert_with_message(quant_conv1.weight_quant is not quant_conv2.weight_quant)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sharing a proxy\n",
    "\n",
    "The second one, which we are introducing now, allows to share the same quantization instance among multiple layers. This is done by simply sharing the proxy wrapping it. This can be useful in those scenarios where, for example, we want different layers to share the same scale factor. The syntax goes as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.446499Z",
     "iopub.status.busy": "2025-05-09T14:48:40.445800Z",
     "iopub.status.idle": "2025-05-09T14:48:40.482418Z",
     "shell.execute_reply": "2025-05-09T14:48:40.479191Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=MySignedBinaryWeightQuantizer)\n",
    "quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=quant_conv1.weight_quant)\n",
    "\n",
    "assert_with_message(quant_conv1.weight_quant is quant_conv2.weight_quant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.493493Z",
     "iopub.status.busy": "2025-05-09T14:48:40.492400Z",
     "iopub.status.idle": "2025-05-09T14:48:40.512061Z",
     "shell.execute_reply": "2025-05-09T14:48:40.509211Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message((quant_conv1.weight_quant.scale() == quant_conv2.weight_quant.scale()).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens in background is that the weight quantizer now has access to both `quant_conv1` and `quant_conv2`. \n",
    "So let's say we want to build a quantizer similar to `ParamFromMaxWeightQuantizer`, but in this case we want the scale factor to be initialized with the average of both weight tensors. When a quantizer has access to multiple parent modules, they are passed in at dependency injection time as a *tuple* under the same name `module` as before. So we can do the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.521988Z",
     "iopub.status.busy": "2025-05-09T14:48:40.520906Z",
     "iopub.status.idle": "2025-05-09T14:48:40.567447Z",
     "shell.execute_reply": "2025-05-09T14:48:40.564997Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "class SharedParamFromMeanWeightQuantizer(MySignedBinaryWeightQuantizer):\n",
    "    \n",
    "    @value\n",
    "    def scaling_init(module):\n",
    "        if isinstance(module, tuple):\n",
    "            return torch.cat((module[0].weight.view(-1), module[1].weight.view(-1))).abs().mean()\n",
    "        else:\n",
    "            return module.weight.abs().mean()\n",
    "        \n",
    "quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=SharedParamFromMeanWeightQuantizer)\n",
    "old_quant_conv1_scale = quant_conv1.weight_quant.scale()\n",
    "quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=quant_conv1.weight_quant)\n",
    "new_quant_conv1_scale = quant_conv1.weight_quant.scale()\n",
    "\n",
    "assert_with_message(not (old_quant_conv1_scale == new_quant_conv1_scale).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.576971Z",
     "iopub.status.busy": "2025-05-09T14:48:40.575537Z",
     "iopub.status.idle": "2025-05-09T14:48:40.589646Z",
     "shell.execute_reply": "2025-05-09T14:48:40.586888Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "assert_with_message((new_quant_conv1_scale == quant_conv2.weight_quant.scale()).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how, when `quant_conv2` is initialized using the `weight_quant` of `quant_conv1`, weight quantization is re-initialized for both layers such that they end up having the same scale.\n",
    "\n",
    "We can see in this example how Brevitas works consistently with Pytorch's eager execution model. When we initialize `quant_conv1` we still don't know that its weight quantizer is going to be shared with `quant_conv2`, and the semantics of Pytorch impose that `quant_conv1` should work correctly both before and after `quant_conv2` is declared. The way we take advantage of dependency injection allows to do so."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sharing an instance of Activation Quantization\n",
    "\n",
    "Sharing an instance of activation quantization is easier because for most scenarios it's enough to simply share the whole layer itself, e.g. calling the same `QuantReLU` from multiple places in the forward pass.\n",
    "\n",
    "For those scenarios where sharing the whole layer is not possible, there is something important to keep in mind. Instances of activation quantization include (for performance reasons) the implementation of the non-linear activation itself (if any). So, for example, using a `QuantReLU.act_quant` to initialize a `QuantConv2d.output_quant` should be avoided as we would not share not only the quantizer, but also the relu activation function.  \n",
    "In general then sharing of instances of activations quantization should be done only between activations of the same *kind*. \n",
    "\n",
    "*Note*: we say kind and not type because `input_quant`, `output_quant` and `IdentityQuant` count as being the same kind of activation, even though they belong to different type of layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dealing with Weight Initialization\n",
    "\n",
    "There is a type of situation that Brevitas cannot deal with automatically. That is, when the initialization of the quantizer depends on the layer to which it is applied (like with the `ParamFromMaxWeightQuantizer` or `SharedParamFromMeanWeightQuantizer` quantizers), but the layer gets modified after it is initialized. \n",
    "\n",
    "The typical example is with weight initialization when training from scratch (so rather than loading from a floating-point state_dict):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.600440Z",
     "iopub.status.busy": "2025-05-09T14:48:40.599438Z",
     "iopub.status.idle": "2025-05-09T14:48:40.625082Z",
     "shell.execute_reply": "2025-05-09T14:48:40.622936Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "quant_conv_w_init = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)\n",
    "torch.nn.init.uniform_(quant_conv_w_init.weight)\n",
    "\n",
    "assert_with_message(not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see how the scale factor is not initialized correctly anymore. In this case we can simply trigger re-initialization of the weight quantizer manually:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.634989Z",
     "iopub.status.busy": "2025-05-09T14:48:40.634167Z",
     "iopub.status.idle": "2025-05-09T14:48:40.650600Z",
     "shell.execute_reply": "2025-05-09T14:48:40.648673Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "quant_conv_w_init.weight_quant.init_tensor_quant()\n",
    "\n",
    "assert_with_message((quant_conv_w_init.weight.abs().max() == quant_conv_w_init.weight_quant.scale()).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note*: because the way weights are initialized is often the same as to how an optimizer performs the weight update step, there are currently no plans to try to perform re-initialization automatically (as it happens e.g. when a `state_dict` is loaded) since it wouldn't be possible to distinguish between the two scenarios."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a Custom Quantization API\n",
    "\n",
    "Finally, let's go through an even more complicated example. We are going to look at a scenario that illustrates the differences between a standard `Injector` (implemented in the dependencies library) and our `ExtendedInjector` extension.\n",
    "\n",
    "Let's say we want to build two quantizers for respectively weights and activations and build a simple API on top of them.\n",
    "In particular, we want to be able to switch between `BinaryQuant` and `ClampedBinaryQuant` (a variant of binary quantization with clamping), and we want to optionally perform *per-channel scaling*.\n",
    "To do so, we are going to implement the controlling logic through a hierarchy of ExtendedInjector, leaving two boolean flags exposed as arguments of the quantizers, with the idea then that the flags can be set through keyword arguments of the respective quantized layers.\n",
    "\n",
    "We can go as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.660862Z",
     "iopub.status.busy": "2025-05-09T14:48:40.659775Z",
     "iopub.status.idle": "2025-05-09T14:48:40.687418Z",
     "shell.execute_reply": "2025-05-09T14:48:40.685085Z"
    }
   },
   "outputs": [],
   "source": [
    "from brevitas.core.quant import ClampedBinaryQuant\n",
    "from brevitas.proxy import WeightQuantProxyFromInjector, ActQuantProxyFromInjector\n",
    "from brevitas.inject import this\n",
    "\n",
    "\n",
    "class CommonQuantizer(ExtendedInjector):\n",
    "    scaling_impl = ParameterScaling\n",
    "    signed=True\n",
    "    \n",
    "    @value\n",
    "    def tensor_quant(is_clamped):\n",
    "        # returning a class to auto-wire from a value function\n",
    "        # wouldn't be allowed in a standard Injector\n",
    "        if is_clamped:\n",
    "            return ClampedBinaryQuant\n",
    "        else:\n",
    "            return BinaryQuant\n",
    "    \n",
    "    @value\n",
    "    def scaling_shape(scaling_per_output_channel):\n",
    "        if scaling_per_output_channel:\n",
    "            # returning this.something from a value function \n",
    "            # wouldn't be allowed in a standard Injector\n",
    "            return this.per_channel_broadcastable_shape\n",
    "        else:\n",
    "            return ()\n",
    "        \n",
    "        \n",
    "class AdvancedWeightQuantizer(CommonQuantizer):\n",
    "    proxy_class = WeightQuantProxyFromInjector\n",
    "        \n",
    "    @value\n",
    "    def per_channel_broadcastable_shape(module):\n",
    "        return (module.weight.shape[0], 1, 1, 1)\n",
    "    \n",
    "    @value\n",
    "    def scaling_init(module, scaling_per_output_channel):\n",
    "        if scaling_per_output_channel:\n",
    "            num_ch = module.weight.shape[0]\n",
    "            return module.weight.abs().view(num_ch, -1).max(dim=1)[0].view(-1, 1, 1, 1)\n",
    "        else:\n",
    "            return module.weight.abs().max()\n",
    "        \n",
    "        \n",
    "class AdvancedActQuantizer(CommonQuantizer):\n",
    "    scaling_init = 0.01\n",
    "    proxy_class = ActQuantProxyFromInjector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a bunch of things going on here to unpack.  \n",
    "\n",
    "The first one is that a `@value` function can return a class to auto-wire and inject, as seen in the definition of `tensor_quant`. This wouldn't normally be possible with a standard `Injector`, but it's possible with an `ExtendedInjector`. This way we can switch between different implementations of `tensor_quant`.\n",
    "\n",
    "The second one is the special object `this`. `this` is already present in the *dependencies* library, and it's used as a way to retrieve attributes of the quantizer from within the quantizer itself. However, normally it wouldn't be possible to return a reference to `this` from a `@value` function. Again this is something that only a `ExtendedInjector` supports, and it allows to chain different attributes in a way such that the chained values are computed only when necessary. \n",
    "\n",
    "Let's see the quantizers applied to a layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.696832Z",
     "iopub.status.busy": "2025-05-09T14:48:40.695972Z",
     "iopub.status.idle": "2025-05-09T14:48:40.733452Z",
     "shell.execute_reply": "2025-05-09T14:48:40.731277Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[ 0.1612, -0.1612, -0.1612],\n",
       "          [-0.1612, -0.1612, -0.1612],\n",
       "          [ 0.1612,  0.1612,  0.1612]],\n",
       "\n",
       "         [[-0.1612,  0.1612, -0.1612],\n",
       "          [-0.1612,  0.1612,  0.1612],\n",
       "          [-0.1612, -0.1612,  0.1612]],\n",
       "\n",
       "         [[-0.1612,  0.1612,  0.1612],\n",
       "          [ 0.1612,  0.1612, -0.1612],\n",
       "          [ 0.1612,  0.1612,  0.1612]]],\n",
       "\n",
       "\n",
       "        [[[ 0.1924,  0.1924,  0.1924],\n",
       "          [-0.1924, -0.1924,  0.1924],\n",
       "          [-0.1924,  0.1924, -0.1924]],\n",
       "\n",
       "         [[ 0.1924, -0.1924,  0.1924],\n",
       "          [ 0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924, -0.1924, -0.1924]],\n",
       "\n",
       "         [[-0.1924, -0.1924,  0.1924],\n",
       "          [ 0.1924, -0.1924, -0.1924],\n",
       "          [ 0.1924, -0.1924,  0.1924]]]], grad_fn=<MulBackward0>), scale=tensor([[[[0.1612]]],\n",
       "\n",
       "\n",
       "        [[[0.1924]]]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "per_channel_quant_conv = QuantConv2d(\n",
    "    3, 2, (3, 3), \n",
    "    weight_quant=AdvancedWeightQuantizer, \n",
    "    weight_is_clamped=False, \n",
    "    weight_scaling_per_output_channel=True)\n",
    "per_channel_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected the weight scale is now a vector. Everything we said so far about quantizers still applies, so for example we can load the floating-point state dict we defined before and observe how it triggers an update of the weight scale:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.742320Z",
     "iopub.status.busy": "2025-05-09T14:48:40.741604Z",
     "iopub.status.idle": "2025-05-09T14:48:40.767078Z",
     "shell.execute_reply": "2025-05-09T14:48:40.764938Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[[[ 0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924]],\n",
       "\n",
       "         [[-0.1924, -0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924, -0.1924],\n",
       "          [ 0.1924,  0.1924,  0.1924]],\n",
       "\n",
       "         [[ 0.1924,  0.1924, -0.1924],\n",
       "          [-0.1924,  0.1924,  0.1924],\n",
       "          [-0.1924,  0.1924,  0.1924]]],\n",
       "\n",
       "\n",
       "        [[[ 0.1899, -0.1899,  0.1899],\n",
       "          [-0.1899,  0.1899, -0.1899],\n",
       "          [ 0.1899,  0.1899,  0.1899]],\n",
       "\n",
       "         [[-0.1899,  0.1899, -0.1899],\n",
       "          [ 0.1899, -0.1899, -0.1899],\n",
       "          [-0.1899,  0.1899,  0.1899]],\n",
       "\n",
       "         [[ 0.1899,  0.1899, -0.1899],\n",
       "          [-0.1899, -0.1899, -0.1899],\n",
       "          [ 0.1899, -0.1899, -0.1899]]]], grad_fn=<MulBackward0>), scale=tensor([[[[0.1924]]],\n",
       "\n",
       "\n",
       "        [[[0.1899]]]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "per_channel_quant_conv.load_state_dict(float_conv.state_dict())\n",
    "per_channel_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case we have a per-channel quantizer, so the original floating-point weight tensor is now quantized per channel. \n",
    "\n",
    "Similarly, we can apply our custom activation quantizer to e.g. a `QuantIdentity` layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.775874Z",
     "iopub.status.busy": "2025-05-09T14:48:40.775160Z",
     "iopub.status.idle": "2025-05-09T14:48:40.798285Z",
     "shell.execute_reply": "2025-05-09T14:48:40.796148Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0100,  0.0100, -0.0100,  0.0100],\n",
       "        [ 0.0100,  0.0100, -0.0100,  0.0100],\n",
       "        [-0.0100,  0.0100, -0.0100, -0.0100],\n",
       "        [ 0.0100, -0.0100, -0.0100, -0.0100]], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.nn import QuantIdentity\n",
    "\n",
    "quant_identity = QuantIdentity(\n",
    "    act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=False)\n",
    "quant_identity(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how `AdvancedActQuantizer` doesn't define a `per_channel_broadcastable_shape`, yet no errors are triggered. This is because `this.per_channel_broadcastable_shape` is required only when `scaling_per_output_channel` is `True`, while in this case `scaling_per_output_channel` is `False`.\n",
    "Let' try to set it to `True` then:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:40.818238Z",
     "iopub.status.busy": "2025-05-09T14:48:40.817508Z",
     "iopub.status.idle": "2025-05-09T14:48:41.538092Z",
     "shell.execute_reply": "2025-05-09T14:48:41.535562Z"
    },
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "DependencyError",
     "evalue": "'AdvancedActQuantizer' can not resolve attribute 'per_channel_broadcastable_shape'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mDependencyError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[35], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mbrevitas\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m QuantIdentity\n\u001b[0;32m----> 3\u001b[0m quant_identity \u001b[38;5;241m=\u001b[39m \u001b[43mQuantIdentity\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43mact_quant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mAdvancedActQuantizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mis_clamped\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mscaling_per_output_channel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_activation.py:113\u001b[0m, in \u001b[0;36mQuantIdentity.__init__\u001b[0;34m(self, act_quant, return_quant_tensor, **kwargs)\u001b[0m\n\u001b[1;32m    108\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m    109\u001b[0m         \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    110\u001b[0m         act_quant: Optional[ActQuantType] \u001b[38;5;241m=\u001b[39m Int8ActPerTensorFloat,\n\u001b[1;32m    111\u001b[0m         return_quant_tensor: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    112\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m--> 113\u001b[0m     \u001b[43mQuantNLAL\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m    114\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    115\u001b[0m \u001b[43m        \u001b[49m\u001b[43minput_quant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    116\u001b[0m \u001b[43m        \u001b[49m\u001b[43mact_impl\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    117\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpassthrough_act\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    118\u001b[0m \u001b[43m        \u001b[49m\u001b[43mact_quant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mact_quant\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    119\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_quant_tensor\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_quant_tensor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    120\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/quant_layer.py:34\u001b[0m, in \u001b[0;36mQuantNonLinearActLayer.__init__\u001b[0;34m(self, act_impl, passthrough_act, input_quant, act_quant, return_quant_tensor, **kwargs)\u001b[0m\n\u001b[1;32m     32\u001b[0m QuantLayerMixin\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, return_quant_tensor)\n\u001b[1;32m     33\u001b[0m QuantInputMixin\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, input_quant, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 34\u001b[0m \u001b[43mQuantNonLinearActMixin\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mact_impl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpassthrough_act\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mact_quant\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/act.py:66\u001b[0m, in \u001b[0;36mQuantNonLinearActMixin.__init__\u001b[0;34m(self, act_impl, passthrough_act, act_quant, act_proxy_prefix, act_kwargs_prefix, **kwargs)\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m     56\u001b[0m         \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     57\u001b[0m         act_impl: Optional[Type[Module]],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     61\u001b[0m         act_kwargs_prefix\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     62\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     63\u001b[0m     prefixed_kwargs \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m     64\u001b[0m         act_kwargs_prefix \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mact_impl\u001b[39m\u001b[38;5;124m'\u001b[39m: act_impl,\n\u001b[1;32m     65\u001b[0m         act_kwargs_prefix \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpassthrough_act\u001b[39m\u001b[38;5;124m'\u001b[39m: passthrough_act}\n\u001b[0;32m---> 66\u001b[0m     \u001b[43mQuantProxyMixin\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m     67\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     68\u001b[0m \u001b[43m        \u001b[49m\u001b[43mquant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mact_quant\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     69\u001b[0m \u001b[43m        \u001b[49m\u001b[43mproxy_prefix\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mact_proxy_prefix\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     70\u001b[0m \u001b[43m        \u001b[49m\u001b[43mkwargs_prefix\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mact_kwargs_prefix\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     71\u001b[0m \u001b[43m        \u001b[49m\u001b[43mproxy_protocol\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mActQuantProxyProtocol\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     72\u001b[0m \u001b[43m        \u001b[49m\u001b[43mnone_quant_injector\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mNoneActQuant\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mprefixed_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     74\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/nn/mixin/base.py:52\u001b[0m, in \u001b[0;36mQuantProxyMixin.__init__\u001b[0;34m(self, quant, proxy_protocol, none_quant_injector, proxy_prefix, kwargs_prefix, **kwargs)\u001b[0m\n\u001b[1;32m     50\u001b[0m     quant_injector \u001b[38;5;241m=\u001b[39m quant\n\u001b[1;32m     51\u001b[0m     quant_injector \u001b[38;5;241m=\u001b[39m quant_injector\u001b[38;5;241m.\u001b[39mlet(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mfilter_kwargs(kwargs_prefix, kwargs))\n\u001b[0;32m---> 52\u001b[0m     quant \u001b[38;5;241m=\u001b[39m \u001b[43mquant_injector\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mproxy_class\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mquant_injector\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     54\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(quant, proxy_protocol):\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/runtime_quant.py:221\u001b[0m, in \u001b[0;36mActQuantProxyFromInjector.__init__\u001b[0;34m(self, quant_layer, quant_injector)\u001b[0m\n\u001b[1;32m    220\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, quant_layer, quant_injector):\n\u001b[0;32m--> 221\u001b[0m     \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mquant_layer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mquant_injector\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    222\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcache_class \u001b[38;5;241m=\u001b[39m _CachedIO\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/runtime_quant.py:94\u001b[0m, in \u001b[0;36mActQuantProxyFromInjectorBase.__init__\u001b[0;34m(self, quant_layer, quant_injector)\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, quant_layer, quant_injector):\n\u001b[0;32m---> 94\u001b[0m     \u001b[43mQuantProxyFromInjector\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mquant_layer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mquant_injector\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     95\u001b[0m     ActQuantProxyProtocol\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m)\n\u001b[1;32m     96\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mis_passthrough_act \u001b[38;5;241m=\u001b[39m _is_passthrough_act(quant_injector)\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/quant_proxy.py:80\u001b[0m, in \u001b[0;36mQuantProxyFromInjector.__init__\u001b[0;34m(self, quant_layer, quant_injector)\u001b[0m\n\u001b[1;32m     78\u001b[0m \u001b[38;5;66;03m# Use a normal list and not a ModuleList since this is a pointer to parent modules\u001b[39;00m\n\u001b[1;32m     79\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtracked_module_list \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m---> 80\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd_tracked_module\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquant_layer\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     81\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdisable_quant \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m     82\u001b[0m \u001b[38;5;66;03m# Torch.compile compatibility requires this\u001b[39;00m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/quant_proxy.py:120\u001b[0m, in \u001b[0;36mQuantProxyFromInjector.add_tracked_module\u001b[0;34m(self, module)\u001b[0m\n\u001b[1;32m    118\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtracked_module_list\u001b[38;5;241m.\u001b[39mappend(module)\n\u001b[1;32m    119\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupdate_tracked_modules()\n\u001b[0;32m--> 120\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minit_tensor_quant\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    121\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    122\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTrying to add None as a parent module.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/proxy/runtime_quant.py:149\u001b[0m, in \u001b[0;36mActQuantProxyFromInjectorBase.init_tensor_quant\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    148\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minit_tensor_quant\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m--> 149\u001b[0m     tensor_quant \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquant_injector\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtensor_quant\u001b[49m\n\u001b[1;32m    150\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mact_impl\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mquant_injector:\n\u001b[1;32m    151\u001b[0m         act_impl \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mquant_injector\u001b[38;5;241m.\u001b[39mact_impl\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/_dependencies/this.py:51\u001b[0m, in \u001b[0;36m_ThisSpec.__call__\u001b[0;34m(self, __self__)\u001b[0m\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m kind \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 51\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mgetattr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msymbol\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     52\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m DependencyError:\n\u001b[1;32m     53\u001b[0m         message \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m     54\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou tried to shift this more times than Injector has levels\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     55\u001b[0m         )\n",
      "File \u001b[0;32m/proj/xlabs/users/nfraser/opt/miniforge3/envs/20231115_brv_pt1.13.1/lib/python3.10/site-packages/brevitas/inject/__init__.py:129\u001b[0m, in \u001b[0;36m_ExtendedInjectorType.__getattr__\u001b[0;34m(cls, attrname)\u001b[0m\n\u001b[1;32m    126\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    127\u001b[0m         message \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{!r}\u001b[39;00m\u001b[38;5;124m can not resolve attribute \u001b[39m\u001b[38;5;132;01m{!r}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m    128\u001b[0m             \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, current_attr)\n\u001b[0;32m--> 129\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m DependencyError(message)\n\u001b[1;32m    131\u001b[0m marker, attribute, args, have_defaults \u001b[38;5;241m=\u001b[39m spec\n\u001b[1;32m    133\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mset\u001b[39m(args)\u001b[38;5;241m.\u001b[39missubset(cached):\n",
      "\u001b[0;31mDependencyError\u001b[0m: 'AdvancedActQuantizer' can not resolve attribute 'per_channel_broadcastable_shape'"
     ]
    }
   ],
   "source": [
    "from brevitas.nn import QuantIdentity\n",
    "\n",
    "quant_identity = QuantIdentity(\n",
    "    act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected we get an error saying that the quantizer cannot resolve `per_channel_broadcastable_shape`. If we pass it in then we can get a per-channel quantizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-05-09T14:48:41.549650Z",
     "iopub.status.busy": "2025-05-09T14:48:41.548924Z",
     "iopub.status.idle": "2025-05-09T14:48:41.578388Z",
     "shell.execute_reply": "2025-05-09T14:48:41.576357Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IntQuantTensor(value=tensor([[ 0.0100,  0.0100,  0.0100, -0.0100],\n",
       "        [ 0.0100, -0.0100, -0.0100,  0.0100],\n",
       "        [-0.0100, -0.0100, -0.0100,  0.0100],\n",
       "        [ 0.0100,  0.0100,  0.0100,  0.0100]], grad_fn=<MulBackward0>), scale=tensor([[0.0100],\n",
       "        [0.0100],\n",
       "        [0.0100],\n",
       "        [0.0100]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quant_identity = QuantIdentity(\n",
    "    act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True,\n",
    "    per_channel_broadcastable_shape=(4, 1), return_quant_tensor=True)\n",
    "quant_identity(torch.randn(4, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have seen how powerful dependency injection is. In a way, it's even too expressive. For users that are not interesting in building completely custom quantizers, it can be hard to make sense of how the various components available under `brevitas.core` can be assembled together according to best practices."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
